import io
import sys

from llama_index.core import Settings
from llama_index.core.agent import ReActAgent
from llama_index.core.tools import FunctionTool

use_ollama = False

if use_ollama:
    from llama_index.llms.ollama import Ollama

    llm = Ollama(model="qwen2.5:7b-instruct-q4_0", request_timeout=120.0)
else:
    from llama_index.llms.openllm import OpenLLM
    import os

    llm = OpenLLM(
        # model="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
        model="Qwen/Qwen2-7B-Instruct",
        api_base="https://api.siliconflow.cn/v1",
        api_key=os.environ.get("SILICON_API_KEY"),
        max_tokens=2048,
        temperature=0.1,
    )
    # from llama_index.llms.deepseek import DeepSeek
    # llm = DeepSeek(model="deepseek-chat")

    # llm = OpenLLM(
    #     model="Qwen/Qwen2.5-7B-Instruct", api_base="http://10.0.100.210/inf-app/a14565593800110080709472/v1",
    #     api_key="123456",
    #     is_chat_model=True,
    #     max_tokens=2048,
    #     temperature=0.1,
    # )

Settings.llm = llm


def execute_python_code(code: str) -> str:
    """执行 Python 代码并捕获输出。请注意，在代码最后必须显式调用“print”将结果输出！"""

    # Create a StringIO object to capture the output
    old_stdout = sys.stdout
    new_stdout = io.StringIO()
    sys.stdout = new_stdout

    try:
        # Using `exec` to execute code
        exec(code)
    except Exception as e:
        # If an exception occurs, capture the exception information
        output = str(e)
    else:
        # If the execution is successful, capture the output
        output = new_stdout.getvalue()
    finally:
        # Recover the standard output
        sys.stdout = old_stdout

    # Wrap the output and status into a ServiceResponse object
    return output


code_exec_tool = FunctionTool.from_defaults(fn=execute_python_code)

# 使用react agent
agent = ReActAgent.from_tools([code_exec_tool], llm=llm, verbose=True)
response = agent.chat("我的电脑是macos系统，首先帮我查找用户名，然后帮我查找我的home目录下的Downloads下最大的文件")
print(str(response))

# code = '''
# def factorial(n):
#     result = 1
#     for i in range(1, n+1):
#         result *= i
#     return result
#
# print(factorial(5))
# '''
# print(execute_python_code(code))